{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using MCP Tools with aisuite\n",
    "\n",
    "This notebook demonstrates how to use MCP (Model Context Protocol) servers with aisuite to give AI models access to external tools and data sources.\n",
    "\n",
    "## What is MCP?\n",
    "\n",
    "MCP (Model Context Protocol) is a standardized protocol that allows AI applications to connect to external data sources and tools. MCP servers expose tools, resources, and prompts that AI models can use.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "Install aisuite with MCP support:\n",
    "```bash\n",
    "pip install 'aisuite[mcp]'\n",
    "# Or install providers you need:\n",
    "pip install 'aisuite[openai,mcp]'\n",
    "```\n",
    "\n",
    "You'll also need to install an MCP server. For this example, we'll use the filesystem server:\n",
    "```bash\n",
    "npm install -g @modelcontextprotocol/server-filesystem\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import aisuite as ai\n",
    "from aisuite.mcp import MCPClient\n",
    "\n",
    "# Load environment variables (API keys)\n",
    "load_dotenv()\n",
    "\n",
    "# Verify API key is set\n",
    "if not os.getenv(\"OPENAI_API_KEY\"):\n",
    "    raise ValueError(\"Please set OPENAI_API_KEY environment variable\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Basic MCP Tool Usage\n",
    "\n",
    "Let's connect to a filesystem MCP server and use it to read files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to the filesystem MCP server\n",
    "# This gives the AI access to files in the specified directory\n",
    "mcp_client = MCPClient(\n",
    "    command=\"npx\",\n",
    "    args=[\"-y\", \"@modelcontextprotocol/server-filesystem\", os.getcwd()]\n",
    ")\n",
    "\n",
    "print(f\"Connected to MCP server: {mcp_client}\")\n",
    "print(f\"\\nAvailable tools:\")\n",
    "for tool in mcp_client.list_tools():\n",
    "    print(f\"  - {tool['name']}: {tool.get('description', 'No description')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create aisuite client\n",
    "client = ai.Client()\n",
    "\n",
    "# Use MCP tools with aisuite\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Please read the README.md file and summarize what this project does.\"}\n",
    "]\n",
    "\n",
    "response = client.chat.completions.create(\n",
    "    model=\"openai:gpt-4o\",\n",
    "    messages=messages,\n",
    "    tools=mcp_client.get_callable_tools(),  # MCP tools work like regular Python functions!\n",
    "    max_turns=3  # Allow multiple tool calls\n",
    ")\n",
    "\n",
    "print(\"\\nAI Response:\")\n",
    "print(response.choices[0].message.content)\n",
    "\n",
    "# View the tool interactions\n",
    "print(\"\\n\" + \"=\"*60)\n",
    "print(\"Tool Call History:\")\n",
    "print(\"=\"*60)\n",
    "for msg in response.choices[0].intermediate_messages:\n",
    "    if msg.role == \"assistant\" and msg.tool_calls:\n",
    "        for tool_call in msg.tool_calls:\n",
    "            print(f\"\\nTool: {tool_call.function.name}\")\n",
    "            print(f\"Arguments: {tool_call.function.arguments}\")\n",
    "    elif msg.role == \"tool\":\n",
    "        print(f\"Result: {msg.content[:200]}...\")  # First 200 chars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Mixing MCP Tools with Python Functions\n",
    "\n",
    "You can seamlessly mix MCP tools with regular Python functions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a custom Python function\n",
    "def get_current_time() -> str:\n",
    "    \"\"\"Get the current date and time.\"\"\"\n",
    "    from datetime import datetime\n",
    "    return datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "\n",
    "def calculate_word_count(text: str) -> int:\n",
    "    \"\"\"Calculate the number of words in a text.\n",
    "    \n",
    "    Args:\n",
    "        text: The text to count words in\n",
    "    \"\"\"\n",
    "    return len(text.split())\n",
    "\n",
    "# Mix MCP tools and Python functions\n",
    "all_tools = mcp_client.get_callable_tools() + [get_current_time, calculate_word_count]\n",
    "\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\", \n",
    "        \"content\": \"What time is it? Also, read the README.md file and tell me how many words it contains.\"\n",
    "    }\n",
    "]\n",
    "\n",
    "response = client.chat.completions.create(\n",
    "    model=\"openai:gpt-4o\",\n",
    "    messages=messages,\n",
    "    tools=all_tools,  # Both MCP and Python tools!\n",
    "    max_turns=5\n",
    ")\n",
    "\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 3: Using Specific MCP Tools\n",
    "\n",
    "You can also select specific tools instead of using all of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get only specific tools\n",
    "read_file = mcp_client.get_tool(\"read_file\")\n",
    "list_directory = mcp_client.get_tool(\"list_directory\")\n",
    "\n",
    "# Use only these tools\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"List all Python files in the current directory.\"}\n",
    "]\n",
    "\n",
    "response = client.chat.completions.create(\n",
    "    model=\"openai:gpt-4o\",\n",
    "    messages=messages,\n",
    "    tools=[read_file, list_directory],  # Only specific tools\n",
    "    max_turns=2\n",
    ")\n",
    "\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 4: Using MCP with Different Providers\n",
    "\n",
    "MCP tools work with all aisuite providers!"
   ]
  },
  {
   "cell_type": "markdown",
   "source": "## Example 4: Using MCP Config Dict Format (Simplified)\n\nInstead of creating an MCPClient explicitly, you can pass MCP server configuration directly as a dict in the `tools` parameter! This is more convenient for simple use cases.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try with different providers\n",
    "providers = [\n",
    "    \"openai:gpt-4o\",\n",
    "    \"anthropic:claude-3-5-sonnet-20240620\",\n",
    "]\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"List the files in the current directory and tell me how many there are.\"}\n",
    "]\n",
    "\n",
    "for model in providers:\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\"Provider: {model}\")\n",
    "    print(f\"{'='*60}\\n\")\n",
    "    \n",
    "    try:\n",
    "        response = client.chat.completions.create(\n",
    "            model=model,\n",
    "            messages=messages,\n",
    "            tools=mcp_client.get_callable_tools(),\n",
    "            max_turns=3\n",
    "        )\n",
    "        print(response.choices[0].message.content)\n",
    "    except Exception as e:\n",
    "        print(f\"Error: {e}\")\n",
    "        print(\"Make sure you have the API key set and provider installed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 5: Connecting to Multiple MCP Servers\n",
    "\n",
    "You can connect to multiple MCP servers and combine their tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to a second MCP server (example - adjust as needed)\n",
    "# For demonstration, we'll create another filesystem server for a different directory\n",
    "\n",
    "# Create a temp directory for demonstration\n",
    "import tempfile\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "\n",
    "# Second MCP client\n",
    "mcp_client_2 = MCPClient(\n",
    "    command=\"npx\",\n",
    "    args=[\"-y\", \"@modelcontextprotocol/server-filesystem\", temp_dir]\n",
    ")\n",
    "\n",
    "# Combine tools from both servers\n",
    "all_mcp_tools = (\n",
    "    mcp_client.get_callable_tools() + \n",
    "    mcp_client_2.get_callable_tools()\n",
    ")\n",
    "\n",
    "print(f\"Total tools available: {len(all_mcp_tools)}\")\n",
    "\n",
    "# Note: In practice, you might want to rename tools or use namespacing\n",
    "# to avoid conflicts between servers with similar tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "It's good practice to close MCP connections when done. You can also use MCPClient as a context manager."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Close the MCP connections\n",
    "mcp_client.close()\n",
    "mcp_client_2.close()\n",
    "\n",
    "print(\"MCP connections closed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Context Manager (Recommended)\n",
    "\n",
    "The recommended way to use MCPClient is with a context manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using context manager ensures proper cleanup\n",
    "with MCPClient(\n",
    "    command=\"npx\",\n",
    "    args=[\"-y\", \"@modelcontextprotocol/server-filesystem\", os.getcwd()]\n",
    ") as mcp:\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"openai:gpt-4o\",\n",
    "        messages=[{\"role\": \"user\", \"content\": \"How many files are in the current directory?\"}],\n",
    "        tools=mcp.get_callable_tools(),\n",
    "        max_turns=2\n",
    "    )\n",
    "    print(response.choices[0].message.content)\n",
    "\n",
    "# Connection is automatically closed after the with block"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Key takeaways:\n",
    "\n",
    "1. **Easy Integration**: MCP tools work seamlessly with aisuite's existing tool system\n",
    "2. **Mix and Match**: Combine MCP tools with regular Python functions\n",
    "3. **Provider Agnostic**: Works with any aisuite provider (OpenAI, Anthropic, Google, etc.)\n",
    "4. **Multiple Servers**: Connect to multiple MCP servers simultaneously\n",
    "5. **Simple API**: Just `MCPClient()` → `get_callable_tools()` → pass to `tools=[]`\n",
    "\n",
    "For more MCP servers, check out:\n",
    "- https://github.com/modelcontextprotocol/servers\n",
    "- Official MCP documentation: https://modelcontextprotocol.io/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}